"""
RAG工具类
为Agent提供检索增强生成功能
"""

import logging
from typing import List, Optional, Dict, Any

from langchain_community.vectorstores import Qdrant
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.llms import Tongyi
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from qdrant_client import QdrantClient

from utils.config import Config

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RAGTool:
    """RAG工具类，提供检索增强生成功能"""
    
    def __init__(self, collection_name: str = None):
        """
        初始化RAG工具
        
        Args:
            collection_name: 向量数据库集合名称
        """
        # 验证配置
        Config.validate_config()
        
        self.collection_name = collection_name or Config.QDRANT_COLLECTION_NAME
        
        # 初始化嵌入模型
        self.embeddings = DashScopeEmbeddings(
            model=Config.EMBEDDING_MODEL,
            dashscope_api_key=Config.DASHSCOPE_API_KEY
        )
        
        # 初始化大语言模型
        self.llm = Tongyi(
            model_name=Config.LLM_MODEL,
            dashscope_api_key=Config.DASHSCOPE_API_KEY,
            temperature=0.1,
            max_tokens=2000
        )
        
        # 初始化Qdrant客户端
        self.qdrant_client = QdrantClient(
            path=Config.QDRANT_PATH
        )
        
        # 初始化向量存储
        self.vector_store = Qdrant(
            client=self.qdrant_client,
            collection_name=self.collection_name,
            embeddings=self.embeddings
        )
        
        # 创建检索器
        self.retriever = self.vector_store.as_retriever(
            search_type="similarity",
            search_kwargs={"k": 5}  # 检索前5个最相关的文档片段
        )
        
        # 创建提示模板
        self.prompt_template = PromptTemplate(
            template="""基于以下上下文信息回答问题。如果上下文中没有相关信息，请回答"根据提供的信息无法回答该问题"。

上下文信息：
{context}

问题：{question}

请基于上下文信息提供准确、详细的回答：""",
            input_variables=["context", "question"]
        )
        
        # 创建检索问答链
        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.retriever,
            return_source_documents=True,
            chain_type_kwargs={"prompt": self.prompt_template}
        )
        
        logger.info(f"RAG工具初始化完成，使用集合: {self.collection_name}")
    
    def search_documents(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """
        搜索相关文档
        
        Args:
            query: 搜索查询
            k: 返回的文档数量
            
        Returns:
            相关文档列表
        """
        try:
            # 使用向量存储进行相似性搜索
            docs = self.vector_store.similarity_search_with_score(query, k=k)
            
            results = []
            for doc, score in docs:
                results.append({
                    "content": doc.page_content,
                    "metadata": doc.metadata,
                    "score": score
                })
            
            logger.info(f"搜索查询 '{query}' 返回 {len(results)} 个结果")
            return results
            
        except Exception as e:
            logger.error(f"搜索文档时出错: {str(e)}")
            return []
    
    def answer_question(self, question: str) -> Dict[str, Any]:
        """
        基于检索的文档回答问题
        
        Args:
            question: 用户问题
            
        Returns:
            包含答案和源文档的字典
        """
        try:
            logger.info(f"正在回答问题: {question}")
            
            # 使用检索问答链获取答案
            result = self.qa_chain({"query": question})
            
            response = {
                "question": question,
                "answer": result["result"],
                "source_documents": []
            }
            
            # 处理源文档信息
            for doc in result["source_documents"]:
                response["source_documents"].append({
                    "content": doc.page_content,
                    "metadata": doc.metadata,
                    "source": doc.metadata.get("source", "未知来源")
                })
            
            logger.info("问题回答完成")
            return response
            
        except Exception as e:
            logger.error(f"回答问题时出错: {str(e)}")
            return {
                "question": question,
                "answer": f"抱歉，处理您的问题时出现错误: {str(e)}",
                "source_documents": []
            }
    
    def get_collection_stats(self) -> Dict[str, Any]:
        """
        获取向量数据库集合统计信息
        
        Returns:
            集合统计信息
        """
        try:
            collection_info = self.qdrant_client.get_collection(self.collection_name)
            return {
                "collection_name": self.collection_name,
                "vectors_count": collection_info.vectors_count,
                "status": collection_info.status,
                "config": {
                    "vector_size": collection_info.config.params.vectors.size,
                    "distance": collection_info.config.params.vectors.distance
                }
            }
        except Exception as e:
            logger.error(f"获取集合统计信息时出错: {str(e)}")
            return {}
    
    def update_retriever_config(self, search_type: str = "similarity", k: int = 5):
        """
        更新检索器配置
        
        Args:
            search_type: 搜索类型
            k: 检索文档数量
        """
        try:
            self.retriever = self.vector_store.as_retriever(
                search_type=search_type,
                search_kwargs={"k": k}
            )
            
            # 重新创建问答链
            self.qa_chain = RetrievalQA.from_chain_type(
                llm=self.llm,
                chain_type="stuff",
                retriever=self.retriever,
                return_source_documents=True,
                chain_type_kwargs={"prompt": self.prompt_template}
            )
            
            logger.info(f"检索器配置已更新: search_type={search_type}, k={k}")
            
        except Exception as e:
            logger.error(f"更新检索器配置时出错: {str(e)}")
    
    def batch_answer_questions(self, questions: List[str]) -> List[Dict[str, Any]]:
        """
        批量回答问题
        
        Args:
            questions: 问题列表
            
        Returns:
            答案列表
        """
        results = []
        for question in questions:
            result = self.answer_question(question)
            results.append(result)
        
        logger.info(f"批量处理了 {len(questions)} 个问题")
        return results
